#!/anaconda3/envs/FEALPy/bin python3.8
# -*- coding: utf-8 -*-
# File: app.py
# Author: Bryan SHEN
# E-mail: m18801919240_3@163.com
# Site: Shanghai, China
# Time: 2024/4/22 13:02
# File-Desp: 构建api链接，传入sentence，获取相应的sentence embedding

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import uvicorn
from utils import response_code
from utils.logger import logger
from src.embedding.sentenc_embedding import Embedding


# Assuming the Embedding class has already been defined/imported
class SentenceList(BaseModel):
    sentences: List[str]


app = FastAPI()

# Initialize the Embedding class with your model path
model_path = '/data_hdd/bryan/social-hotspot-collect-and-cluster/models/bce-embedding-base_v1'
embedder = Embedding(model_path)


@app.post("/bce_embeddings/")
async def create_embeddings(sentence_list: SentenceList):
    try:
        # Get embeddings as a torch tensor
        embeddings = embedder.get_embedding(sentence_list.sentences)

        # Convert torch tensor to list for JSON serialization
        embeddings_list = embeddings.tolist()
        return {"embeddings": embeddings_list}
    except Exception as e:
        return response_code.resp_500()

if __name__ == '__main__':

    HOST = '0.0.0.0'
    PORT = 8091
    logger.info(f"Clean Panel server started on port {HOST}:{PORT}")
    uvicorn.run('app:app', host=HOST, port=PORT)